// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/common/testutil/optional_empty.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>

using namespace kumo::pollux::exec;
using namespace kumo::pollux::exec::test;
using namespace kumo::pollux::functions::aggregate::test;

namespace kumo::pollux::aggregate::test {

namespace {

class MapUnionSumTest : public AggregationTestBase {};

TEST_F(MapUnionSumTest, global) {
  auto data = make_row_vector({
      make_nullable_map_vector<int64_t, int64_t>({
          common::testutil::optionalEmpty, // empty map
          std::nullopt, // null map
          {{{1, 10}, {2, 20}}},
          {{{1, 11}, {3, 30}, {4, 40}}},
          {{{3, 30}, {5, 50}, {1, 12}}},
      }),
  });

  auto expected = make_row_vector({
      make_map_vector<int64_t, int64_t>({
          {{1, 33}, {2, 20}, {3, 60}, {4, 40}, {5, 50}},
      }),
  });

  testAggregations({data}, {}, {"map_union_sum(c0)"}, {expected});
}

TEST_F(MapUnionSumTest, globalVarcharKey) {
  std::vector<std::string> keyStrings = {
      "Tall mountains",
      "Wide rivers",
      "Deep oceans",
      "Thick dark forests",
      "Expansive vistas",
  };
  std::vector<StringView> keys;
  for (const auto& key : keyStrings) {
    keys.push_back(StringView(key));
  }

  auto data = make_row_vector({
      make_nullable_map_vector<StringView, int64_t>({
          common::testutil::optionalEmpty, // empty map
          std::nullopt, // null map
          {{{keys[0], 10}, {keys[1], 20}}},
          {{{keys[0], 11}, {keys[2], 30}, {keys[3], 40}}},
          {{{keys[2], 30}, {keys[4], 50}, {keys[0], 12}}},
      }),
  });

  auto expected = make_row_vector({
      make_map_vector<StringView, int64_t>({
          {{keys[0], 33},
           {keys[1], 20},
           {keys[2], 60},
           {keys[3], 40},
           {keys[4], 50}},
      }),
  });

  testAggregations({data}, {}, {"map_union_sum(c0)"}, {expected});
}

TEST_F(MapUnionSumTest, nullAndEmptyMaps) {
  auto allEmptyMaps = make_row_vector({
      make_map_vector<int64_t, int64_t>({
          {},
          {},
          {},
      }),
  });

  auto expectedEmpty = make_row_vector({
      make_map_vector<int64_t, int64_t>({
          {},
      }),
  });

  testAggregations({allEmptyMaps}, {}, {"map_union_sum(c0)"}, {expectedEmpty});

  auto allNullMaps = make_row_vector({
      make_nullable_map_vector<int64_t, int64_t>({
          std::nullopt,
          std::nullopt,
          std::nullopt,
      }),
  });

  auto expectedNull = make_row_vector({
      make_nullable_map_vector<int64_t, int64_t>({
          std::nullopt,
      }),
  });

  testAggregations({allNullMaps}, {}, {"map_union_sum(c0)"}, {expectedNull});

  auto emptyAndNullMaps = make_row_vector({
      make_nullable_map_vector<int64_t, int64_t>({
          std::nullopt,
          common::testutil::optionalEmpty,
          std::nullopt,
          common::testutil::optionalEmpty,
      }),
  });

  testAggregations(
      {emptyAndNullMaps}, {}, {"map_union_sum(c0)"}, {expectedEmpty});
}

TEST_F(MapUnionSumTest, tinyintOverflow) {
  auto data = make_row_vector({
      make_nullable_map_vector<int64_t, int8_t>({
          {{{1, 10}, {2, 20}}},
          {{{1, 100}, {3, 30}, {4, 40}}},
          {{{3, 30}, {5, 50}, {1, 30}}},
      }),
  });

  auto plan = PlanBuilder()
                  .values({data})
                  .singleAggregation({}, {"map_union_sum(c0)"})
                  .planNode();

  POLLUX_ASSERT_THROW(
      AssertQueryBuilder(plan).copyResults(pool()), "Value 140 exceeds 127");

  data = make_row_vector({
      make_nullable_map_vector<int64_t, int8_t>({
          {{{1, -10}, {2, -20}}},
          {{{1, -100}, {3, -30}, {4, -40}}},
          {{{3, -30}, {5, -50}, {1, -30}}},
      }),
  });

  plan = PlanBuilder()
             .values({data})
             .singleAggregation({}, {"map_union_sum(c0)"})
             .planNode();

  POLLUX_ASSERT_THROW(
      AssertQueryBuilder(plan).copyResults(pool()),
      "Value -140 is less than -128");
}

TEST_F(MapUnionSumTest, smallintOverflow) {
  const int16_t largeValue = std::numeric_limits<int16_t>::max() - 20;
  const int16_t smallValue = std::numeric_limits<int16_t>::min() + 20;
  auto data = make_row_vector({
      make_nullable_map_vector<int64_t, int16_t>({
          {{{1, 10}, {2, 20}}},
          {{{1, largeValue}, {3, 30}, {4, 40}}},
          {{{3, 30}, {5, 50}, {1, 30}}},
      }),
  });

  auto plan = PlanBuilder()
                  .values({data})
                  .singleAggregation({}, {"map_union_sum(c0)"})
                  .planNode();

  POLLUX_ASSERT_THROW(
      AssertQueryBuilder(plan).copyResults(pool()),
      "Value 32787 exceeds 32767");

  data = make_row_vector({
      make_nullable_map_vector<int64_t, int16_t>({
          {{{1, -10}, {2, -20}}},
          {{{1, smallValue}, {3, -30}, {4, -40}}},
          {{{3, -30}, {5, -50}, {1, -30}}},
      }),
  });

  plan = PlanBuilder()
             .values({data})
             .singleAggregation({}, {"map_union_sum(c0)"})
             .planNode();

  POLLUX_ASSERT_THROW(
      AssertQueryBuilder(plan).copyResults(pool()),
      "Value -32788 is less than -32768");
}

TEST_F(MapUnionSumTest, integerOverflow) {
  const int32_t largeValue = std::numeric_limits<int32_t>::max() - 20;
  const int32_t smallValue = std::numeric_limits<int32_t>::min() + 20;
  auto data = make_row_vector({
      make_nullable_map_vector<int64_t, int32_t>({
          {{{1, 10}, {2, 20}}},
          {{{1, largeValue}, {3, 30}, {4, 40}}},
          {{{3, 30}, {5, 50}, {1, 30}}},
      }),
  });

  auto plan = PlanBuilder()
                  .values({data})
                  .singleAggregation({}, {"map_union_sum(c0)"})
                  .planNode();

  POLLUX_ASSERT_THROW(
      AssertQueryBuilder(plan).copyResults(pool()),
      "Value 2147483667 exceeds 2147483647");

  data = make_row_vector({
      make_nullable_map_vector<int64_t, int32_t>({
          {{{1, -10}, {2, -20}}},
          {{{1, smallValue}, {3, -30}, {4, -40}}},
          {{{3, -30}, {5, -50}, {1, -30}}},
      }),
  });

  plan = PlanBuilder()
             .values({data})
             .singleAggregation({}, {"map_union_sum(c0)"})
             .planNode();

  POLLUX_ASSERT_THROW(
      AssertQueryBuilder(plan).copyResults(pool()),
      "Value -2147483668 is less than -2147483648");
}

TEST_F(MapUnionSumTest, bigintOverflow) {
  const int64_t largeValue = std::numeric_limits<int64_t>::max() - 20;
  const int64_t smallValue = std::numeric_limits<int64_t>::min() + 20;
  auto data = make_row_vector({
      make_nullable_map_vector<int64_t, int64_t>({
          {{{1, 10}, {2, 20}}},
          {{{1, largeValue}, {3, 30}, {4, 40}}},
          {{{3, 30}, {5, 50}, {1, 30}}},
      }),
  });

  auto plan = PlanBuilder()
                  .values({data})
                  .singleAggregation({}, {"map_union_sum(c0)"})
                  .planNode();

  POLLUX_ASSERT_THROW(
      AssertQueryBuilder(plan).copyResults(pool()),
      "Value 9223372036854775827 exceeds 9223372036854775807");

  data = make_row_vector({
      make_nullable_map_vector<int64_t, int64_t>({
          {{{1, -10}, {2, -20}}},
          {{{1, smallValue}, {3, -30}, {4, -40}}},
          {{{3, -30}, {5, -50}, {1, -30}}},
      }),
  });

  plan = PlanBuilder()
             .values({data})
             .singleAggregation({}, {"map_union_sum(c0)"})
             .planNode();

  POLLUX_ASSERT_THROW(
      AssertQueryBuilder(plan).copyResults(pool()),
      "Value -9223372036854775828 is less than -9223372036854775808");
}

TEST_F(MapUnionSumTest, floatNan) {
  constexpr float kInf = std::numeric_limits<float>::infinity();
  constexpr float kNan = std::numeric_limits<float>::quiet_NaN();

  auto data = make_row_vector({
      make_nullable_map_vector<int64_t, float>({
          {{{1, 10}, {2, 20}}},
          {{{1, kNan}, {3, 30}, {5, 50}}},
          {{{3, 30}, {5, kInf}, {1, 30}}},
      }),
  });

  auto expected = make_row_vector({
      make_map_vector<int64_t, float>({
          {{1, kNan}, {2, 20}, {3, 60}, {5, kInf}},
      }),
  });

  testAggregations({data}, {}, {"map_union_sum(c0)"}, {expected});
}

TEST_F(MapUnionSumTest, doubleNan) {
  constexpr float kInf = std::numeric_limits<double>::infinity();
  constexpr float kNan = std::numeric_limits<double>::quiet_NaN();

  auto data = make_row_vector({
      make_nullable_map_vector<int64_t, double>({
          {{{1, 10}, {2, 20}}},
          {{{1, kNan}, {3, 30}, {5, 50}}},
          {{{3, 30}, {5, kInf}, {1, 30}}},
      }),
  });

  auto expected = make_row_vector({
      make_map_vector<int64_t, double>({
          {{1, kNan}, {2, 20}, {3, 60}, {5, kInf}},
      }),
  });

  testAggregations({data}, {}, {"map_union_sum(c0)"}, {expected});
}

TEST_F(MapUnionSumTest, groupBy) {
  auto data = make_row_vector({
      make_flat_vector<int64_t>({1, 2, 1, 2, 1}),
      make_nullable_map_vector<int64_t, int64_t>({
          {}, // empty map
          std::nullopt, // null map
          {{{1, 10}, {2, 20}}},
          {{{1, 11}, {3, 30}, {4, 40}}},
          {{{3, 30}, {5, 50}, {1, 12}}},
      }),
  });

  auto expected = make_row_vector({
      make_flat_vector<int64_t>({1, 2}),
      make_map_vector<int64_t, int64_t>({
          {{1, 22}, {2, 20}, {3, 30}, {5, 50}},
          {{1, 11}, {3, 30}, {4, 40}},
      }),
  });

  testAggregations({data}, {"c0"}, {"map_union_sum(c1)"}, {expected});
}

TEST_F(MapUnionSumTest, groupByVarcharKey) {
  std::vector<std::string> keyStrings = {
      "Tall mountains",
      "Wide rivers",
      "Deep oceans",
      "Thick dark forests",
      "Expansive vistas",
  };
  std::vector<StringView> keys;
  for (const auto& key : keyStrings) {
    keys.push_back(StringView(key));
  }

  auto data = make_row_vector({
      make_flat_vector<int64_t>({1, 2, 1, 2, 1}),
      make_nullable_map_vector<StringView, int64_t>({
          {}, // empty map
          std::nullopt, // null map
          {{{keys[0], 10}, {keys[1], 20}}},
          {{{keys[0], 11}, {keys[2], 30}, {keys[3], 40}}},
          {{{keys[2], 30}, {keys[4], 50}, {keys[0], 12}}},
      }),
  });

  auto expected = make_row_vector({
      make_flat_vector<int64_t>({1, 2}),
      make_map_vector<StringView, int64_t>({
          {{keys[0], 22}, {keys[1], 20}, {keys[2], 30}, {keys[4], 50}},
          {{keys[0], 11}, {keys[2], 30}, {keys[3], 40}},
      }),
  });

  testAggregations({data}, {"c0"}, {"map_union_sum(c1)"}, {expected});
}

TEST_F(MapUnionSumTest, groupByJsonKey) {
  std::vector<std::string> jsonKeyStrings = {
      "\"key1\"",
      "\"key2\"",
      "\"key3\"",
      "\"key4\"",
      "\"key5\"",
  };
  std::vector<StringView> jsonKeys;
  jsonKeys.reserve(jsonKeyStrings.size());
  for (const auto& key : jsonKeyStrings) {
    jsonKeys.emplace_back(key);
  }
  auto data = make_row_vector({
      make_flat_vector<int64_t>({1, 2, 1, 2, 1}),
      make_nullable_map_vector<StringView, int64_t>({
          {},
          std::nullopt,
          {{{jsonKeys[0], 10}, {jsonKeys[1], 20}}},
          {{{jsonKeys[0], 11}, {jsonKeys[2], 30}, {jsonKeys[3], 40}}},
          {{{jsonKeys[2], 30}, {jsonKeys[4], 50}, {jsonKeys[0], 12}}},
      }),
  });

  auto expected = make_row_vector({
      make_flat_vector<int64_t>({1, 2}),
      make_map_vector<StringView, int64_t>({
          {{jsonKeys[0], 22},
           {jsonKeys[1], 20},
           {jsonKeys[2], 30},
           {jsonKeys[4], 50}},
          {{jsonKeys[0], 11}, {jsonKeys[2], 30}, {jsonKeys[3], 40}},
      }),
  });

  testAggregations({data}, {"c0"}, {"map_union_sum(c1)"}, {expected});
}

TEST_F(MapUnionSumTest, groupByBooleanKeys) {
  auto data = make_row_vector({
      make_flat_vector<int64_t>({1, 2, 1, 2, 1}),
      make_nullable_map_vector<bool, int64_t>({
          {}, // empty map
          std::nullopt, // null map
          {{{true, 10}, {false, 20}}},
          {{{true, 11}, {false, 30}, {true, 40}}},
          {{{false, 28}, {true, 50}, {true, 12}}},
      }),

  });

  auto expected = make_row_vector({
      make_flat_vector<int64_t>({1, 2}),
      make_map_vector<bool, int64_t>({
          {{{true, 72}, {false, 48}}},
          {{{true, 51}, {false, 30}}},
      }),
  });

  testAggregations({data}, {"c0"}, {"map_union_sum(c1)"}, {expected});
}

TEST_F(MapUnionSumTest, floatingPointKeys) {
  auto data = make_row_vector({
      make_flat_vector<int32_t>({1, 2, 1, 2, 1, 1, 2, 2}),
      make_map_vector_from_json<float, int64_t>({
          "{1.1: 10, 1.2: 20, 1.3: 30}",
          "{2.1: 10, 1.2: 20, 2.3: 30}",
          "{3.1: 10, 1.2: 20, 2.3: 30}",
          "{}",
          "null",
          "{4.1: 10, 4.2: 20, 2.3: 30}",
          "{5.1: 10, 4.2: 20, 2.3: 30}",
          "{6.1: 10, 6.2: 20, 6.3: 30}",
      }),
  });

  auto expected = make_row_vector({
      make_map_vector_from_json<float, int64_t>({
          "{1.1: 10, 1.2: 60, 1.3: 30, 2.1: 10, 2.3: 120, 3.1: 10, 4.1: 10, "
          "4.2: 40, 5.1: 10, 6.1: 10, 6.2: 20, 6.3: 30}",
      }),
  });

  testAggregations({data}, {}, {"map_union_sum(c1)"}, {expected});
}

TEST_F(MapUnionSumTest, nanKeys) {
  // Verify that NaNs with different binary representations are considered
  // equal and deduplicated when used as keys in the output map.
  constexpr double kNan = std::numeric_limits<double>::quiet_NaN();
  constexpr double kSNaN = std::numeric_limits<double>::signaling_NaN();

  // Global aggregation.
  auto data = make_row_vector(
      {make_nullable_map_vector<double, int32_t>({
           {{{kSNaN, 10}, {2, 20}}},
           {{{kNan, 1}, {3, 30}, {5, 50}}},
           {{{3, 30}, {kSNaN, 4}, {1, 30}}},
       }),
       make_flat_vector<int32_t>({1, 1, 2})});

  auto expected = make_row_vector({
      make_map_vector<double, int32_t>({
          {{1, 30}, {2, 20}, {3, 60}, {5, 50}, {kNan, 15}},
      }),
  });

  testAggregations({data}, {}, {"map_union_sum(c0)"}, {expected});

  // Group by aggregation.
  expected = make_row_vector(
      {make_map_vector<double, int32_t>({
           {{2, 20}, {3, 30}, {5, 50}, {kNan, 11}},
           {{1, 30}, {3, 30}, {kNan, 4}},
       }),
       make_flat_vector<int32_t>({1, 2})});

  testAggregations(
      {data}, {"c1"}, {"map_union_sum(c0)"}, {"a0", "c1"}, {expected});
}

TEST_F(MapUnionSumTest, complexType) {
  // Verify that NaNs with different binary representations are considered equal
  // and deduplicated when used as keys in the output map.
  static const auto kNaN = std::numeric_limits<double>::quiet_NaN();
  static const auto kSNaN = std::numeric_limits<double>::signaling_NaN();

  // Global Aggregation, Complex type(Row)
  // The complex input values are:
  // [{"key":[1,1],"value":1},{"key":["NaN",2],"value":2},{"key":[2,4],"value":3},{"key":[3,5],"value":4},
  // {"key":["NaN",2],"value":5}, {"key":["NaN",2],"value":6}]
  auto data = make_row_vector(
      {make_map_vector(
           {0, 1, 2, 3, 4, 5},
           make_row_vector(
               {make_flat_vector<double>({1, kSNaN, 2, 3, kNaN, kSNaN}),
                make_flat_vector<int32_t>({1, 2, 4, 5, 2, 2})}),
           make_flat_vector<int32_t>({1, 2, 3, 4, 5, 6})),
       make_flat_vector<int32_t>({1, 1, 1, 2, 2, 2})});

  // The expected result is
  // [{"key":[1,1],"value":1},{"key":[2,4],"value":3},{"key":[3,5],"value":4},
  // {"key":["NaN",2],"value":12}]
  auto expectedResult = make_row_vector({make_map_vector(
      {0},
      make_row_vector(
          {make_flat_vector<double>({1, 2, 3, kNaN}),
           make_flat_vector<int32_t>({1, 4, 5, 2})}),
      make_flat_vector<int32_t>({1, 3, 4, 13}))});

  testAggregations({data}, {}, {"map_union_sum(c0)"}, {expectedResult});

  // Group by Aggregation, Complex type(Row)
  // The expected result is
  // [{"key":[1,1],"value":1},{"key":[2,4],"value":3},
  //  {"key":["NaN",2],"value":2}] | 1
  // [{"key":[3,5],"value":4},{"key":["NaN",2],"value":11}] | 2
  expectedResult = make_row_vector(
      {make_map_vector(
           {0, 3},
           make_row_vector(
               {make_flat_vector<double>({1, 2, kNaN, 3, kNaN}),
                make_flat_vector<int32_t>({1, 4, 2, 5, 2})}),
           make_flat_vector<int32_t>({1, 3, 2, 4, 11})),
       make_flat_vector<int32_t>({1, 2})});

  testAggregations(
      {data}, {"c1"}, {"map_union_sum(c0)"}, {"a0", "c1"}, {expectedResult});
}

TEST_F(MapUnionSumTest, unknownKey) {
  auto data = make_row_vector({make_all_null_map_vector(3, UNKNOWN(), BIGINT())});
  testAggregations({data}, {}, {"map_union_sum(c0)"}, {"VALUES (NULL)"});
}

TEST_F(MapUnionSumTest, decimalKey) {
  // test on nulls
  auto null_data =
      make_row_vector({make_all_null_map_vector(3, DECIMAL(10, 5), BIGINT())});
  testAggregations({null_data}, {}, {"map_union_sum(c0)"}, {"VALUES (NULL)"});

  // test on non-null decimal keys
  auto data = make_map_vector<int128_t, int64_t>(
      {{{{1000}, 2}, {{1001}, 1}}, {{{1000}, 1}, {{1001}, 1}}},
      MAP(DECIMAL(30, 2), BIGINT()));
  auto expected = make_row_vector({make_map_vector<int128_t, int64_t>(
      {{{{1000}, 3}, {{1001}, 2}}}, MAP(DECIMAL(30, 2), BIGINT()))});
  testAggregations(
      {make_row_vector({data})}, {}, {"map_union_sum(c0)"}, {expected});
}

} // namespace
} // namespace kumo::pollux::aggregate::test
